def get_img_output_length(width, height):
    """计算VGG16特征提取后的特征图展平长度"""
    def get_output_length(input_length):
        filter_sizes = [2, 2, 2, 2, 2]  # 各池化层核大小
        stride = 2
        for i in range(len(filter_sizes)):
            input_length = (input_length - filter_sizes[i]) // stride + 1
        return input_length
    return get_output_length(width) * get_output_length(height)

class SiameseNetwork(nn.Module):
    def __init__(self, input_shape):
        super(SiameseNetwork, self).__init__()
        self.vgg = vgg16.features
        
        # 计算展平后的特征向量长度
        flat_shape = 512 * get_img_output_length(input_shape[1], input_shape[0])
        
        # 定义全连接分类器
        self.fc = nn.Sequential(
            nn.Linear(flat_shape, 512),
            nn.ReLU(inplace=True),
            nn.Linear(512, 256),
            nn.ReLU(inplace=True),
            nn.Linear(256, 1)  # 二分类输出
        )
        
    def forward_once(self, x):
        output = self.vgg(x)
        output = torch.flatten(output, 1)
        return output
        
    def forward(self, input1, input2):
        output1 = self.forward_once(input1)
        output2 = self.forward_once(input2)
        # 计算特征差异并分类
        output = torch.abs(output1 - output2)  # 使用绝对值差异
        output = self.fc(output)
        return output
